{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import oemof.solph as solph\n",
    "import oemof.graph as graph\n",
    "import networkx as nx\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restore an energysystem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energysystem = solph.EnergySystem()\n",
    "energysystem.restore(dpath=None, filename=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energysystem_graph = graph.create_nx_graph(energysystem)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### You can pickle graph to save it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_file_name = 'energysystem_graph'\n",
    "nx.readwrite.write_gpickle(G=energysystem_graph, path=graph_file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define function for drawing graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_graph(grph, edge_labels=True, node_color='#AFAFAF',\n",
    "               edge_color='#CFCFCF', plot=True, node_size=2000,\n",
    "               with_labels=True, arrows=True, layout='neato'):\n",
    "    \"\"\"\n",
    "    Draw a graph. This function will be removed in future versions.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    grph : networkxGraph\n",
    "        A graph to draw.\n",
    "    edge_labels : boolean\n",
    "        Use nominal values of flow as edge label\n",
    "    node_color : dict or string\n",
    "        Hex color code oder matplotlib color for each node. If string, all\n",
    "        colors are the same.\n",
    "\n",
    "    edge_color : string\n",
    "        Hex color code oder matplotlib color for edge color.\n",
    "\n",
    "    plot : boolean\n",
    "        Show matplotlib plot.\n",
    "\n",
    "    node_size : integer\n",
    "        Size of nodes.\n",
    "\n",
    "    with_labels : boolean\n",
    "        Draw node labels.\n",
    "\n",
    "    arrows : boolean\n",
    "        Draw arrows on directed edges. Works only if an optimization_model has\n",
    "        been passed.\n",
    "    layout : string\n",
    "        networkx graph layout, one of: neato, dot, twopi, circo, fdp, sfdp.\n",
    "    \"\"\"\n",
    "    if type(node_color) is dict:\n",
    "        node_color = [node_color.get(g, '#AFAFAF') for g in grph.nodes()]\n",
    "\n",
    "    # set drawing options\n",
    "    options = {\n",
    "     'prog': 'dot',\n",
    "     'with_labels': with_labels,\n",
    "     'node_color': node_color,\n",
    "     'edge_color': edge_color,\n",
    "     'node_size': node_size,\n",
    "     'arrows': arrows\n",
    "    }\n",
    "\n",
    "    # draw graph\n",
    "    pos = nx.drawing.nx_agraph.graphviz_layout(grph, prog=layout)\n",
    "\n",
    "    nx.draw(grph, pos=pos, **options)\n",
    "\n",
    "    # add edge labels for all edges\n",
    "    if edge_labels is True and plt:\n",
    "        labels = nx.get_edge_attributes(grph, 'weight')\n",
    "        nx.draw_networkx_edge_labels(grph, pos=pos, edge_labels=labels)\n",
    "\n",
    "    # show output\n",
    "    if plot is True:\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read back pickled graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energysystem_graph = nx.readwrite.read_gpickle('energysystem_graph')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Draw graph\n",
    "\n",
    "To define colors, use color palette generators, e.g.:   \n",
    "http://javier.xyz/cohesive-colors/  \n",
    "https://colourco.de/  \n",
    "http://seaborn.pydata.org/tutorial/color_palettes.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matplotlib.rcParams['figure.figsize'] = [10.0, 10.0]\n",
    "draw_graph(energysystem_graph, node_size=5500,\n",
    "            node_color={\n",
    "                         'coal': '#0f2e2e',\n",
    "                         'gas': '#c76c56',\n",
    "                         'oil': '#494a19',\n",
    "                         'lignite': '#56201d',\n",
    "                         'bel': '#9a9da1',\n",
    "                         'bth': '#cd3333',\n",
    "                         'wind': '#4ca7c3',\n",
    "                         'pv': '#ffde32',\n",
    "                         'demand_el': '#9a9da1',\n",
    "                         'excess_el': '#9a9da1',\n",
    "                         'demand_th': '#cd3333',\n",
    "                         'pp_coal': '#0f2e2e',\n",
    "                         'pp_lig': '#56201d',\n",
    "                         'pp_gas': '#c76c56',\n",
    "                         'pp_oil': '#494a19',\n",
    "                         'pp_chp': '#eeac7e',\n",
    "                         'b_heat_source': '#cd3333',\n",
    "                         'heat_source': '#cd3333',\n",
    "                         'heat_pump': '#42c77a'},\n",
    "            edge_color='#eeac7e')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
